import logging
import sys
# from sqlalchemy import text
# from sqlalchemy import create_engine
import csv
import pymysql
import re
import requests
import string
import random
import pandas as pd
import os
import datetime
import json
import base64
import os
import time
from alibabacloud_tea_openapi.models import Config
from alibabacloud_searchplat20240529.client import Client
from alibabacloud_searchplat20240529.models import (
    CreateImageAnalyzeTaskRequestDocument,
    CreateImageAnalyzeTaskRequest,
    CreateImageAnalyzeTaskResponse,
    GetImageAnalyzeTaskStatusRequest,
    GetImageAnalyzeTaskStatusResponse
)





# 更新数据表中的一条数据
def update_data(task_id, data_path, pre_anno_path, anno_path, status='成功'):
    conn = pymysql.connect(
        host='9.134.192.120',
        # port=3861,
        user='root',
        password='1987720sgf',
        database='knowledge',
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

    # # 配置数据库连接信息
    # DATABASE_URI = "mysql+mysqldb://root:1987720sgf@9.134.192.120/knowledge?charset=utf8mb4"
    #
    # # 创建一个数据库引擎
    # engine = create_engine(DATABASE_URI)

    with conn.cursor() as connection:
        sql = f"UPDATE data_pipeline SET `status` = '{status}', `data_path` = '{data_path}', `pre_anno_path` = '{pre_anno_path}', `anno_path` = '{anno_path}' WHERE id={task_id}"

        print(sql)
        result = connection.execute(sql)
        conn.commit()


def path2base64(img_path):
    import base64
    with open(img_path, 'rb') as f:
        image_data = f.read()
    base64_str = str(base64.b64encode(image_data), 'utf-8')
    return base64_str


def is_image_only_string(s):
    image_pattern = r'$$.*?$$$.*?$'
    return bool(re.fullmatch(image_pattern, s))


def extract_text_from_string(s):
    # 先匹配图片链接部分
    image_pattern = r'$$.*?$$$.*?$'
    image_match = re.search(image_pattern, s)

    if image_match:
        # 如果找到图片链接，提取图片链接之后的部分作为文字
        text_start_index = image_match.end()
        text = s[text_start_index:].strip()
    else:
        # 如果没有找到图片链接，整个字符串就是文字
        text = s.strip()

    return text


def type_cls(image_paths, description):
    prompt = """你是一个可以根据图片内容提炼相关问题的专业工作人员，综合分析用户上传的图片（image）和对该图片的描述（description）判断出图片的场景类别。并提炼符合场景类型的问题和对应答案。最终返回场景类别，分类原因，提炼出的问题和答案。
    【场景判断步骤】
    1. 解析用户上传的1张或者多张图片，分析内容。
    2. 分析description中的内容。
    3. 结合图片分析结果，description 和【场景判定指导原则】进行分类。
    4. 根据场景类型，参照提炼【问题提炼原则】提炼符合场景的问题和答案
    【场景判定指导原则】
    - 数据图理解: 该场景用户上传的图片是数据图(折线，饼图，柱状图，散点图等等)，也可以根据description的内容判断是否是数据图
    - 单/多表格理解: 该场景用户上传的图片是单个或者多个表格，也可以根据description的内容判断
    - 混合图像理解: 该场景用户上传的图片可能既有表格，也有数据图，或者一个数据图中有多个形式，比如又有折线图又有柱状图。也可以根据description的内容判断
    - 文字信息提取: 图片中含有明显的文字内容。也可以根据description的内容判断。 注意：如果字体过小，则不建议归为此类型。
    - 测量读数:图片中通常是某些测量工具（比如尺子）。也可以根据description的内容判断
    - 指针型读数:图片中通常是某些 表/时钟/电压表/转速表等目标。也可以根据description的内容判断
    - 地图解析:用户上传的图片往往是一个类似地图的图片。也可以根据description的内容判断
    - 文档阅读理解: 用户上传的图片是某个文档中的截图，比如word文档，用户手册，pdf截图等。也可以根据description的内容判断
    - 广告理解: 通常图片是海报/广告文案/电商平台中的商品广告等场景。也可以根据description的内容判断
    - 聊天截图问答: 图片是聊天内容的截图。也可以根据description的内容判断
    - 手写/印刷区分: 图片中的文字部分有手写的内容。也可以根据description的内容判断
    - 行业知识储备：图片内容通常涉及到比较专业的领域，比如医疗/法律/泛金融/文旅/教育/饮食/科技/游戏娱乐/出行/传媒&泛互/工业/智慧零售/运营商&通信/影音娱乐/美食知识储备/穿搭美妆知识储备/ 以及其他行业场景。也可以根据description的内容判断。希望能返回具体的行业场景
    - 语种理解：图片内容通常有除了中文外的其他国家的文字内容。也可以根据description的内容判断
    - 证件：图片内容往往是银行卡/身份证/营业执照等证件。也可以根据description的内容判断
    - 内容描述: 如果不属于以上任何场景，则为此场景
    
    【问题提炼原则】
    - 数据图理解: 主要询问将数据数据图中的全部内容或部分内容转换成表格， 比如：将6点到18点的数据制作成表格
    - 单/多表格理解: 主要询问需要将1个或多个表格进行计算才能得到的内容，比如：第三季度的销售额是多少？谁的成绩最低？价格不高的能力项有哪些？租宝马7系、本田CRV、卡罗拉，这3辆车，价位分别都是多少？
    - 混合图像理解: 与 单/多表格理解 场景的提问方法一致。
    - 文字信息提取: 统一提问：请将图片中的文字内容识别给我，区分大小写。
    - 测量读数: 主要询问图片中测量工具当前的读书是多少。 比如：厚度是多少？ 测量出的数值是多少？ 物品宽度是多少？图中管子的直径是多少？ 
    - 指针型读数: 主要询问图片中的读数是多少。 比如：这两个手表的时间相差多少？当前显示的是几点？图中的压力值读数是多少？
    - 地图解析： 主要对着地图询问指路型的问题。比如：地图中有几个游客中心？我所在的位置靠近起点吗？我现在的位置在哪里？地图上有几个入口可以进去？
    - 文档阅读理解: 结合文档中的语义提炼相关问题和答案。比如：写出第四段作者想表达什么？妈妈的账单里包括哪些内容？天线位置在什么地方？
    - 广告理解: 结合图片内容提炼问题，比如编写广告文案或者要求提炼广告内容或者识别广告中的要素，描述广告中的产品特点等
    - 聊天截图问答：结合聊天内容提炼适当的问题。 比如：刘静想表达什么？ 图中的“我”因为什么被要求退群？图中表示要给谁发红包？每个红包包多少？有什么含义
    - 行业知识储备：请集合各行各业的专业知识，提炼相关的问题。比如：给我些投资的建议？血小板计数是否正常范围内？分析下这两个单子，此人视力怎么样，需要配镜吗？
    - 证件：主要询问识别证件中的关键信息， 比如：身份证号， 银行卡号
    - 语种理解： 要求讲图片中的文字进行翻译。
    - 手写/印刷区分: 要求输出图片中手写的部分。
    - 内容描述: 统一询问： 请描述当这张图片
    【examples】
    {
        "description": "《条形统计图和折线统计图》课件ppt", 
        "verification": {"type": "数据图理解", "reason": "图中包含统计图和折线图，description也说明含有统计图和折线图，所以是数据图理解", "question": "第二季度的销售额是多少", "answer": "55w" } 
    }
    {
        "description": "商务平纹笔记本 10本礼品定制可印logo企业创意广告品定制", 
        "verification": {"type": "广告理解", "reason": "描述中提到广告品定制，所以是广告理解场景" , "question": "请为图中的产品生成广告词", "answer": "xx品牌笔记本，打工人最好的选择"} 
    }
    【注意】输出严格按下面的格式进行输出
    
    {
        "type": "广告理解", 
        "reason": "描述中提到广告品定制，所以是广告理解场景" , 
        "question": "请为图中的产品生成广告词", 
        "answer": "xx品牌笔记本，打工人最好的选择"
    } 
    【下面是真实的description, 请给出场景分类，分类原因，提炼的问题和答案】
    description:
    """ + description

    from zhipuai import ZhipuAI

    img_path = image_paths[0]
    with open(img_path, 'rb') as img_file:
        img_base = base64.b64encode(img_file.read()).decode('utf-8')

    client = ZhipuAI(api_key="5ffdcf14622669c64487ab73392af88a.YkcYrzsq4prA0Yhm")  # 填写您自己的APIKey
    response = client.chat.completions.create(
        model="glm-4v-plus",  # 填写需要调用的模型名称
        messages=[
            {
                "role": "user",  # user , system, assitants(助手)
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_base
                        }
                    },
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
    )
    print(response.choices[0].message.content)


    result = response.choices[0].message.content
    result = result.replace("```json", "").replace("```", "").strip()

    return result


def generate_random_string(length):
    characters = string.ascii_letters + string.digits
    random_string = ''.join(random.choice(characters) for _ in range(length))
    return random_string


def generate_filename(dir_name='mllm_single_wenzi_data'):
    # 获取当前日期
    current_date = datetime.datetime.now().strftime("%Y%m%d")

    # 生成4个随机字符
    random_chars = ''.join(random.choices(string.ascii_letters + string.digits, k=4))

    # 拼接文件名
    filename = f"{dir_name}_{current_date}_{random_chars}"

    return filename


if __name__ == "__main__":
    # # 获取命令行参数
    # args = sys.argv[1:]
    # csv_path = args[0]
    base_dir = '/Users/cainsun/Downloads/wukong_test'
    img_dir = os.path.join(base_dir, 'wukong_test')
    csv_path = os.path.join(base_dir, 'wukong_test.csv')

    columns = ['image_name', 'type', 'reason', 'question', 'answer', 'desc']
    result_df = pd.DataFrame(columns=columns)

    df = pd.read_csv(csv_path)
    i = 0
    for index, row in df.iterrows():
        image_name = row['dir']
        desc = row['text']
        image_list = [os.path.join(img_dir, image_name)]
        result = None
        try:
            result = json.loads(type_cls(image_list, desc))
            if result['type'] != '内容描述':
                result['image_name'] = image_name
                result['desc'] = desc
                result_df.loc[len(result_df)] = result
                print(result)

            i += 1
            if i > 50:
                break
        except Exception as e:
            print(f"报错：{e}, 模型返回:{result}")
    result_df.to_csv(os.path.join(base_dir, 'data.csv'), index=False)


